#include "physics.h"
#include <cmath>

static inline double sqr(double x){return x*x;}

Deriv eval(const State& s, const PhysParams& p){
    Deriv d{};
    for(int i=0;i<3;++i){
        d.dp[i]=s.b[i].v;
        d.dv[i]={0.0,0.0};
    }
    for(int i=0;i<3;++i){
        for(int j=i+1;j<3;++j){
            Vec2 r = s.b[j].p - s.b[i].p;
            double d2 = sqr(r.x)+sqr(r.y)+sqr(p.eps);
            double inv = 1.0/std::sqrt(d2);
            double inv3 = inv*inv*inv;
            Vec2 f = r * (p.G * s.b[i].m * s.b[j].m * inv3);
            d.dv[i] = d.dv[i] + (f / s.b[i].m);
            d.dv[j] = d.dv[j] - (f / s.b[j].m);
        }
    }
    return d;
}

void integrate_rk4(State& s, const PhysParams& p){
    State s1=s;
    Deriv k1=eval(s1,p);
    State s2=s;
    for(int i=0;i<3;++i){
        s2.b[i].p = s.b[i].p + k1.dp[i]*(p.dt*0.5);
        s2.b[i].v = s.b[i].v + k1.dv[i]*(p.dt*0.5);
    }
    Deriv k2=eval(s2,p);
    State s3=s;
    for(int i=0;i<3;++i){
        s3.b[i].p = s.b[i].p + k2.dp[i]*(p.dt*0.5);
        s3.b[i].v = s.b[i].v + k2.dv[i]*(p.dt*0.5);
    }
    Deriv k3=eval(s3,p);
    State s4=s;
    for(int i=0;i<3;++i){
        s4.b[i].p = s.b[i].p + k3.dp[i]*p.dt;
        s4.b[i].v = s.b[i].v + k3.dv[i]*p.dt;
    }
    Deriv k4=eval(s4,p);
    for(int i=0;i<3;++i){
        s.b[i].p = s.b[i].p + (k1.dp[i] + k2.dp[i]*2.0 + k3.dp[i]*2.0 + k4.dp[i])*(p.dt/6.0);
        s.b[i].v = s.b[i].v + (k1.dv[i] + k2.dv[i]*2.0 + k3.dv[i]*2.0 + k4.dv[i])*(p.dt/6.0);
    }
}

double total_energy(const State& s, const PhysParams& p){
    double ke=0.0; for(int i=0;i<3;++i){ ke += 0.5*s.b[i].m*(sqr(s.b[i].v.x)+sqr(s.b[i].v.y)); }
    double pe=0.0; for(int i=0;i<3;++i) for(int j=i+1;j<3;++j){
        Vec2 r = s.b[j].p - s.b[i].p; double d2 = sqr(r.x)+sqr(r.y)+sqr(p.eps); pe -= p.G*s.b[i].m*s.b[j].m/std::sqrt(d2);
    }
    return ke+pe;
}

void reset_initial(State& s){
    reset_figure_eight(s);
}

void reset_figure_eight(State& s){
    // equal-mass periodic figure-eight initial conditions (scaled)
    s.b[0].m = 1.0; s.b[1].m = 1.0; s.b[2].m = 1.0;
    s.b[0].p = { -0.97000436,  0.24308753 };
    s.b[1].p = {  0.0       ,  0.0        };
    s.b[2].p = {  0.97000436, -0.24308753 };
    s.b[0].v = {  0.4662036850,  0.4323657300 };
    s.b[1].v = { -0.93240737 , -0.86473146  };
    s.b[2].v = {  0.4662036850,  0.4323657300 };
}

void reset_solar(State& s){
    // heavy central mass with two light bodies in near-circular orbits
    double M0=5.0; double m1=0.01; double m2=0.02; double r1=1.5; double r2=2.5; double G=1.0;
    s.b[0].m=M0; s.b[1].m=m1; s.b[2].m=m2;
    s.b[0].p={0.0,0.0}; s.b[0].v={0.0,0.0};
    s.b[1].p={ r1, 0.0}; s.b[2].p={-r2,0.0};
    double v1=std::sqrt(G*M0/r1); double v2=std::sqrt(G*M0/r2);
    s.b[1].v={0.0, v1}; s.b[2].v={0.0,-v2};
}

void reset_figure_eight_loose(State& s){
    reset_figure_eight(s);
    double th = 0.05; double cs=std::cos(th), sn=std::sin(th);
    double sp=1.03, sv=0.97;
    for(int i=0;i<3;++i){
        double x = s.b[i].p.x, y = s.b[i].p.y; s.b[i].p.x = sp*(cs*x - sn*y); s.b[i].p.y = sp*(sn*x + cs*y);
        double vx = s.b[i].v.x, vy = s.b[i].v.y; s.b[i].v.x = sv*(cs*vx - sn*vy); s.b[i].v.y = sv*(sn*vx + cs*vy);
    }
}